import argparse
import random
from itertools import chain
import os
from pathlib import Path

def gpower(seed, mat_size):
   out = []
   for row,col in seed:
      r_off = row*mat_size[0]
      c_off = col*mat_size[1]
      for ele in seed:
         r = ele[0] + r_off
         c = ele[1] + c_off
         out.append((r,c))
   return out, (mat_size[0]**2,mat_size[1]**2)

if __name__ == '__main__':
   parser = argparse.ArgumentParser()
   parser.add_argument('--seed_file', default=None, help='path to the seed matrix file')
   parser.add_argument('--matrix_param', default=None, help='random graph parameter (row,col,nnz_per_row)')
   parser.add_argument('--power', type=int, default=2, help='generated power of the seed, defualt is 2')
   parser.add_argument('--output', default=None, help='output file name')

   args = parser.parse_args()
   
   # at list one mode is selected
   assert (not args.seed_file is None) or (not args.matrix_param is None)
   
   seed_mat = list()
   seed_size = None

   out = None
   out_size = None

   if not args.seed_file is None:
      with open(args.seed_file, 'r') as f:
         header = f.readline().split()
         seed_size = (int(header[0]), int(header[1]))
         line = f.readline()
         while line:
            coord = [int(x) for x in line.split()]
            seed_mat.append(tuple(coord[0:2]))
            print("ele: ({},{})".format(coord[0], coord[1]))
            line = f.readline()
         f.close()

      pmat = seed_mat
      psize = seed_size
      for i in range(1, args.power):
         pmat, psize = gpower(pmat, psize)
      
      out,out_size = pmat,psize
   
   elif not args.matrix_param is None:
      params = [int(x) for x in args.matrix_param.split(",")]
      out_size = params[0:2]
      out = []
      for row in range(out_size[0]):
         cols = random.sample(list(chain(range(0,row), range(row+1, out_size[1]))), params[2])
         # cols = random.sample(list(range(row+1, min(row+10,out_size[1]))), params[2])
         for edge in [(row,x) for x in cols]:
            out.append(edge)

   if args.output is None:
      of = args.seed_file.strip('.mtx')+'_power_{}.mtx'.format(args.power)
   elif not args.output is None:
      of = args.output
   else:
      of = 'default_output.mtx'

   with open(of, 'w') as f:
      f.write("%%MatrixMarket matrix coordinate pattern general\n")
      f.write("{} {} {}\n".format(out_size[0], out_size[1], len(out)))
      for r,c in out:
         f.write("{} {}\n".format(r,c))
      f.close()
   

   
